#pragma once
#include <zmq.h>
#include <string>
#include <unordered_map>
#include <functional>
#include <memory>
#include <atomic>
#include <mutex>
#include <thread>
#include <unistd.h>
#include "pzmq_data.h"

// 补充定义 RPC 关键字
#define ZMQ_RPC_FUN (ZMQ_REP | 0x80)
#define ZMQ_RPC_CALL (ZMQ_REQ | 0x80)

class pzmq
{
public:
    using rpc_callback_fun = std::function<std::string(pzmq *, const std::shared_ptr<pzmq_data> &)>;
    using msg_callback_fun = std::function<void(pzmq *, const std::shared_ptr<pzmq_data> &)>;

private:
    const int _rpc_url_head_length = 6;
    std::string _rpc_url_head = "ipc:///tmp/rpc.";

    void *_zmq_ctx;    // ZeroMQ 上下文
    void *_zmq_socket; // ZeroMQ 套接字

    std::unordered_map<std::string, rpc_callback_fun> _zmq_fun; // RPC 回调函数映射表
    std::mutex _zmq_fun_mtx;                                    // 互斥锁
    std::atomic<bool> _flag;                                    // 线程控制标志
    std::unique_ptr<std::thread> _zmq_thread;                   // 事件循环线程

    std::string _rpc_server; // RPC 服务器名称

    int _mode;            // 当前模式
    std::string _zmq_url; // 当前连接的url
    int _timeout;         // 超时时间(毫秒)

public:
    // 指定 RPC 服务器名称（RPC通信创建-惰性）
    pzmq(const std::string &server)
        : _zmq_ctx(nullptr), _zmq_socket(nullptr), _flag(true), _rpc_server(server), _timeout(3000)
    {
        if (server.find("://") != std::string::npos)
        {
            _rpc_url_head.clear();
        }
    }

    // 具体通信模式创建
    pzmq(const std::string &url, int mode, const msg_callback_fun &raw_call = nullptr)
    {
        if ((url[0] != 'i') && (url[1] != 'p'))
        {
            _rpc_url_head.clear();
        }

        if ((_mode != ZMQ_RPC_FUN) || (_mode != ZMQ_RPC_CALL))
            Create(url, raw_call);
    }

    ~pzmq()
    {
        if (_zmq_socket)
        {
            _flag = true;
            zmq_ctx_shutdown(_zmq_ctx);
            if (_zmq_thread)
            {
                _zmq_thread->join();
            }
            CloseZmq();
        }
    }

    // 注册RPC动作
    int RegisterRpcAction(const std::string &action, const rpc_callback_fun &raw_call)
    {
        int ret = 0;
        std::unique_lock<std::mutex> lock(_zmq_fun_mtx);
        // 更新
        if (_zmq_fun.find(action) != _zmq_fun.end())
        {
            _zmq_fun[action] = raw_call;
            return ret;
        }

        // 第一次注册会创建信道，并且创建第一个rpc服务 action_list
        if (_zmq_fun.empty())
        {
            std::string url = _rpc_url_head + _rpc_server;
            _mode = ZMQ_RPC_FUN;
            _zmq_fun["action_list"] = std::bind(&pzmq::ShowRpcActionList, this);
            ret = Create(url);
        }
        _zmq_fun[action] = raw_call;
        return ret;
    }

    // 注销RPC动作
    void UnregisterRpcAction(const std::string &action)
    {
        std::unique_lock<std::mutex> lock(_zmq_fun_mtx);
        if (_zmq_fun.find(action) != _zmq_fun.end())
        {
            _zmq_fun.erase(action);
        }
    }

    // 调用RPC动作
    int CallRpcAction(const std::string &action, const std::string &data, const msg_callback_fun &raw_call)
    {
        int ret = 0;
        std::shared_ptr<pzmq_data> msg_ptr = std::make_shared<pzmq_data>();
        try
        {
            if (_zmq_socket == nullptr)
            {
                if (_rpc_server.empty())
                    return -1;

                std::string url = _rpc_url_head + _rpc_server;
                _mode = ZMQ_RPC_CALL;
                ret = Create(url);
                if (ret)
                    throw ret;
            }

            // 发送action
            zmq_send(_zmq_socket, action.c_str(), action.size(), ZMQ_SNDMORE);
            // 发送data
            zmq_send(_zmq_socket, data.c_str(), data.size(), 0);

            // 接收数据
            zmq_msg_recv(msg_ptr->Get(), _zmq_socket, 0);

            // 执行回调函数
            raw_call(this, msg_ptr);
        }
        catch (int e)
        {
            return e;
        }
        return ret;
    }

    // 列出所有RPC动作
    std::string ShowRpcActionList()
    {
        std::string action_list;
        action_list.reserve(128);
        action_list += "{\"actions\":[";
        for (auto it : _zmq_fun)
        {
            action_list += "\"";
            action_list += it.first;
            action_list += "\", ";
        }
        if (action_list.substr(action_list.size() - 2) == ", ")
        {
            action_list.pop_back();
            action_list.pop_back();
        }
        action_list += "]}";
        return action_list;
    }

    int Create(const std::string &url, const msg_callback_fun &raw_call = nullptr)
    {
        _zmq_url = url;
        do
        {
            _zmq_ctx = zmq_ctx_new();
        } while (_zmq_ctx == nullptr);

        do
        {
            _zmq_socket = zmq_socket(_zmq_ctx, _mode & 0x3f);
        } while (_zmq_socket == nullptr);

        switch (_mode)
        {
        case ZMQ_PUB:
            return CreatePub(url);
            break;
        case ZMQ_SUB:
            return CreateSub(url, raw_call);
            break;
        case ZMQ_PULL:
            return CreatePull(url, raw_call);
            break;
        case ZMQ_PUSH:
            return CreatePush(url);
            break;
        case ZMQ_RPC_FUN:
            return CreateRep(url, raw_call);
            break;
        case ZMQ_RPC_CALL:
            return CreateReq(url);
            break;
        default:
            break;
        }
        return 0;
    }

    int CreatePub(const std::string &url)
    {
        return zmq_bind(_zmq_socket, url.c_str());
    }

    int CreateSub(const std::string &url, const msg_callback_fun &raw_call)
    {
        // 首次重连间隔：100ms
        int reconnect_ivl = 100;
        zmq_setsockopt(_zmq_socket, ZMQ_RECONNECT_IVL, &reconnect_ivl, sizeof(reconnect_ivl));

        // 最大重试间隔：1秒
        int max_reconnect_ivl = 1000;
        zmq_setsockopt(_zmq_socket, ZMQ_RECONNECT_IVL_MAX, &max_reconnect_ivl, sizeof(max_reconnect_ivl));

        int ret = zmq_connect(_zmq_socket, _zmq_url.c_str());

        // 接受所有消息
        zmq_setsockopt(_zmq_socket, ZMQ_SUBSCRIBE, "", 0);
        _flag = false;

        // 开启线程进行处理
        _zmq_thread = std::make_unique<std::thread>(std::bind(&pzmq::ZmqEventLoop, this, raw_call));
        return ret;
    }

    int CreatePush(const std::string &url)
    {
        // 首次重连间隔：100ms
        int reconnect_ivl = 100;
        zmq_setsockopt(_zmq_socket, ZMQ_RECONNECT_IVL, &reconnect_ivl, sizeof(reconnect_ivl));

        // 最大重试间隔：1秒
        int max_reconnect_ivl = 1000;
        zmq_setsockopt(_zmq_socket, ZMQ_RECONNECT_IVL_MAX, &max_reconnect_ivl, sizeof(max_reconnect_ivl));

        zmq_setsockopt(_zmq_socket, ZMQ_SNDTIMEO, &_timeout, sizeof(_timeout));

        return zmq_connect(_zmq_socket, _zmq_url.c_str());
    }
    int CreatePull(const std::string &url, const msg_callback_fun &raw_call)
    {
        int ret = zmq_bind(_zmq_socket, url.c_str());
        _flag = false;

        // 开启线程进行处理
        _zmq_thread = std::make_unique<std::thread>(std::bind(&pzmq::ZmqEventLoop, this, raw_call));
        return ret;
    }

    int CreateReq(const std::string &url)
    {
        if (!_rpc_url_head.empty())
        {
            std::string socket_file = url.substr(_rpc_url_head_length);
            // 判断请求资源是否存在
            if (access(socket_file.c_str(), F_OK) != 0)
            {
                return -1;
            }
        }

        // 发送/接收数据阻塞等待时间
        zmq_setsockopt(_zmq_socket, ZMQ_SNDTIMEO, &_timeout, sizeof(_timeout));
        zmq_setsockopt(_zmq_socket, ZMQ_RCVTIMEO, &_timeout, sizeof(_timeout));
        return zmq_connect(_zmq_socket, url.c_str());
    }
    int CreateRep(const std::string &url, const msg_callback_fun &raw_call)
    {
        int ret = zmq_bind(_zmq_socket, url.c_str());
        _flag = false;

        // 开启线程进行处理
        _zmq_thread = std::make_unique<std::thread>(std::bind(&pzmq::ZmqEventLoop, this, raw_call));
        return ret;
    }

    void ZmqEventLoop(const msg_callback_fun &raw_call)
    {
        // 设置线程名
        pthread_setname_np(pthread_self(), "zmq_event_loop");

        zmq_pollitem_t items[1];
        int ret = 0;
        if (_mode == ZMQ_PULL)
        {
            items[0].socket = _zmq_socket;
            items[0].events = ZMQ_POLLIN; // 只关注读事件
            items[0].revents = 0;
        }

        while (!_flag.load())
        {
            std::shared_ptr<pzmq_data> msg_ptr = std::make_shared<pzmq_data>();
            if (_mode == ZMQ_PULL)
            {
                ret = zmq_poll(items, 1, -1);
                if (ret == -1)
                {
                    zmq_close(_zmq_socket);
                    continue;
                }

                if (!(items[0].revents && ZMQ_POLLIN))
                {
                    continue;
                }
            }

            // 三种模式统一在这接收数据
            // rpc只是接收到第一帧数据：action（服务名字）
            ret = zmq_msg_recv(msg_ptr->Get(), _zmq_socket, 0);
            if (ret <= 0)
            {
                msg_ptr.reset();
                continue;
            }

            if (_mode == ZMQ_RPC_FUN)
            {
                // rpc接收第二帧数据
                std::shared_ptr<pzmq_data> msg_ptr1 = std::make_shared<pzmq_data>();
                ret = zmq_msg_recv(msg_ptr1->Get(), _zmq_socket, 0);
                std::string retval;
                try
                {
                    std::unique_lock<std::mutex> lock(_zmq_fun_mtx);
                    retval = _zmq_fun[msg_ptr->String()](this, msg_ptr1);
                }
                catch (...)
                {
                    retval = "NoAction";
                }
                zmq_send(_zmq_socket, retval.c_str(), retval.size(), 0);
                msg_ptr1.reset();
            }
            else
            {
                // 其他通信方式调用传入的回调函数
                raw_call(this, msg_ptr);
            }
            msg_ptr.reset();
        }
    }

    int send_data(const std::string &raw)
    {
        return zmq_send(_zmq_socket, raw.c_str(), raw.size(), 0);
    }

    void CloseZmq()
    {
        zmq_close(_zmq_socket);
        zmq_ctx_destroy(_zmq_ctx);
        if ((_mode == ZMQ_PUB) || (_mode == ZMQ_PULL) || (_mode == ZMQ_RPC_FUN))
        {
            if (!_rpc_url_head.empty())
            {
                std::string socket_file = _rpc_url_head.substr(_rpc_url_head_length);
                if (access(socket_file.c_str(), F_OK) == 0)
                {
                    remove(socket_file.c_str());
                }
            }
        }
        _zmq_socket = nullptr;
        _zmq_ctx = nullptr;
    }
};